package cn.amebius.simpledocker.websocket;

import cn.amebius.simpledocker.config.DockerFactory;
import cn.amebius.simpledocker.entity.model.DockerClientModel;
import cn.amebius.simpledocker.exception.WebsocketErrorException;
import com.Ostermiller.util.CircularByteBuffer;
import com.github.dockerjava.api.DockerClient;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 功能描述：基础Websocket 抽象类
 *
 * @author HeJiaWei
 * @date 2/6/24
 */
@Slf4j
@Component
public class BaseAbstractWebsocket extends TextWebSocketHandler {

    private String socketName;

    private static DockerFactory dockerFactory;

    // input输入流
    private static final Map<String, InputStream> ins = new ConcurrentHashMap<>();

    // output输出流
    private static final Map<String, OutputStream> outs = new ConcurrentHashMap<>();

    private static final Map<String,WebSocketSession> connections = new ConcurrentHashMap<>();

    private static final AtomicInteger poolSize = new AtomicInteger();
    private static final String DEFAULT_NAME = "default";

    // 无参构造
    public BaseAbstractWebsocket() {
        this.socketName = DEFAULT_NAME;
        initInfos();
    }

    // 有参构造
    public BaseAbstractWebsocket(String name) {
        this.socketName = name;
        initInfos();
    }

    @Autowired
    public void setService(DockerFactory dockerFactory) {
        BaseAbstractWebsocket.dockerFactory = dockerFactory;
    }

    /**
     * 功能描述：websocket线程开启后，IO的初始化操作
     *         生成唯一的input和output流，存入map中
     *
     * @author HeJiaWei
     * @date 2/6/24
     */
    public void initIOs(String sessionId) {
        if (null == sessionId || ins.containsKey(sessionId) || outs.containsKey(sessionId)) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：initIOs参数错误");
        }
        final CircularByteBuffer cbb = new CircularByteBuffer();
        final InputStream in = cbb.getInputStream();
        final OutputStream out = cbb.getOutputStream();
        ins.put(sessionId, in);
        outs.put(sessionId, out);
    }

    /**
     * 功能描述：websocket线程开启后，session池的加入操作
     *
     * @author HeJiaWei
     * @date 2/6/24
     */
    public void initConnections(WebSocketSession session) {
        if (null == session) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：initConnections参数错误");
        }
        String id = session.getId();
        if (ins.containsKey(id) || outs.containsKey(id)) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：initConnections参数错误");
        }
        connections.put(session.getId(),session);
        log.info("Websocket连接池：当前大小[{}],session连接数量为： {}",poolSize.get(),connections.size());
    }

    /**
     * 功能描述： session池子大小
     *
     * @author HeJiaWei
     * @date 2/6/24
     */
    public Integer getConnectionSize() {
        return connections.size();
    }

    public WebSocketSession getConnect(String sessionId) {
        if (null == sessionId || connections.containsKey(sessionId)) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：getConnect参数错误");
        }
        return connections.get(sessionId);
    }

    /**
     * 功能描述：获取输入流
     *
     * @author HeJiaWei
     * @date 2/6/24
     */
    public InputStream getInput(String sessionId) {
        if (null == sessionId) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：initIOs参数错误");
        }
        return ins.get(sessionId);
    }

    public DockerFactory getDockerFactory() {
        return this.dockerFactory;
    }

    public DockerClient getDockerClient(WebSocketSession session) {
        if (null == session) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：getDockerClient参数错误");
        }
        Integer key = (Integer) session.getAttributes().get("key");
        DockerClientModel dcm = getDockerFactory().getByKey(key);
        return dcm.getDockerClient();
    }

    /**
     * 功能描述：获取输出流
     *
     * @author HeJiaWei
     * @date 2/6/24
     */
    public OutputStream getOutput(String sessionId) {
        if (null == sessionId) {
            throw new WebsocketErrorException("BaseAbstractWebsocket：initIOs参数错误");
        }
        return outs.get(sessionId);
    }

    public void initInfos() {
        int size = poolSize.incrementAndGet();
        log.info("Websocket线程池：当前大小为[{}], 初始化 {} 完成",size,this.socketName);
    }

    public void destory(WebSocketSession session) {
        try {
            String id = session.getId();
            session.close();
            InputStream is = getInput(id);
            OutputStream os = getOutput(id);
            if (null != is) {
                is.close();
            }
            if (null != os) {
                os.close();
            }
            ins.remove(id);
            outs.remove(id);
            connections.remove(id);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
